#include <iostream>
#include <iomanip>
#include <cstring>
#include <fstream>
#include <sstream>

#include "tinyxml.h"

#include "TileModule.h"

#define TILES_PTRS		0x2F125B	// 4 byte SNES pointer table
#define ARRS_PTRS		0x2F12AB	// 4 byte SNES pointer table

using namespace std;

void Subpalette::readFromRom(uint addr, Rom& rom) {
  int i;
  uint col;
  for (i=0; i<16; i++) {
    col = rom.readMulti(addr + 2*i, 2) & 0x7fff;
    c[i].r = ((col & 0x001f) << 3);
    c[i].g = (((col & 0x03e0) >> 5) << 3);
    c[i].b = (((col >> 10) & 0x1f) << 3);
  }
}

void Subpalette::writeToRom(uint addr, Rom& rom) {
  int i;
  uint colorBlock;
  for (i=0; i<16; i++) {
    colorBlock = (c[i].r >> 3) & 0x1f;
    colorBlock |= ((c[i].g >> 3) & 0x1f) << 5;
    colorBlock |= ((c[i].b >> 3) & 0x1f) << 10;
    colorBlock &= 0x7fff;

    // Dirty hack to preserve color0 data
    // colorBlock |= (rom.readMulti(addr + 2*i, 2) & 0x8000);

    rom.writeMulti(colorBlock, addr + 2*i, 2);
  }
}

void Subpalette::getAsString(char* out) {
  int i;
  for (i=0; i<16; i++) {
    itoa(c[i].r >> 3, out++, 32);
    itoa(c[i].g >> 3, out++, 32);
    itoa(c[i].b >> 3, out++, 32);
  }
}

void Subpalette::setAsString(char* str) {
  int i;
  for (i=0; i<16; i++) {
    c[i].r = chtoi(*(str++)) << 3;
    c[i].g = chtoi(*(str++)) << 3;
    c[i].b = chtoi(*(str++)) << 3;
  }
}

Palette::Palette(int mt, int mp) {
  mtset = mt;
  mpal = mp;
}

Palette::Palette(char* str) {
  setAsString(str);
}

int Palette::getMapTileset() {
  return mtset;
}

int Palette::getMapPalette() {
  return mpal;
}

void Palette::readFromRom(uint start, Rom& rom) {
  int i;
//  cout << "Reading " << mtset << "/" << mpal << " @ 0x" << hex << start << dec << endl;
  tmpAddr = start - 0x1A0200;
  flag = rom.readMulti(start, 2);
  tmpAltAddr = rom.readMulti(start + 0x20, 2);
  spritePal = rom.read(start + 0x40);
  flashEffect = rom.read(start + 0x60);
  for (i=0; i<6; i++)
    subpals[i].readFromRom(start + i*32, rom);
}

void Palette::writeToRom(uint addr, Rom& rom) {
  int i;
  for (i=0; i<6; i++)
    subpals[i].writeToRom(addr + i*32, rom);
  tmpAddr = addr - 0x1A0200;
  rom.writeMulti(flag, addr, 2);
  rom.write(spritePal, addr + 0x40);
  rom.write(flashEffect, addr + 0x60);
}

uint Palette::getAddr() {
  return tmpAddr;
}

uint Palette::getAltAddr() {
  return tmpAltAddr;
}

uint Palette::getFlag() {
  return flag;
}

void Palette::setFlag(uint f) {
  flag = f;
}

void Palette::setAltPal(int mt, int mp) {
  alt_mtset = mt;
  alt_mpal = mp;
}

int Palette::getAltTset() {
  return alt_mtset;
}

int Palette::getAltPal() {
  return alt_mpal;
}

uchar Palette::getSpritePalette() {
  return spritePal;
}

void Palette::setSpritePalette(uchar sp) {
  spritePal = sp;
}

uchar Palette::getFlashEffect() {
  return flashEffect;
}

void Palette::setFlashEffect(uchar fe) {
  flashEffect = fe;
}

void Palette::getAsString(char* out) {
  int i;
  itoa(mtset, out, 32);
  itoa(mpal, out+1, 32);
  out +=2;

  for (i=0; i<6; i++) {
    subpals[i].getAsString(out);
    out += 16*3;
  }
}

void Palette::setAsString(char* str) {
  int i;
  mtset = chtoi(str[0]);
  mpal = chtoi(str[1]);

  for (i=0; i<6; i++)
    subpals[i].setAsString(str+2+(i*16*3));
}

void Tileset::clearPalettes() {
  list<Palette*>::iterator it;
  for (it = pals.begin(); it != pals.end(); it++)
    delete (*it);
  pals.clear();
}

std::list<Palette*>* Tileset::palettesList() {
  return &pals;
}

void Tileset::setTileAsStr(int tn, char* str) {
  int i,j;
  for (i=0; i<8; i++)
    for(j=0; j<8; j++)
      tiles[tn][i][j] = (uchar) ((str[i*8+j] | 32) % 39 - 9);
}

void Tileset::getTileAsStr(int tn, char* out) {
  if (tn >= NUM_TILES)
    strcpy(out,"0000000000000000000000000000000000000000000000000000000000000000");
  else {
    int i, j;
    for (i=0; i<8; i++) {
      for (j=0; j<8; j++) {
        if (tiles[tn][i][j] <= 9)
          out[i*8+j] = ((char) tiles[tn][i][j]) + 0x30;
        else
          out[i*8+j] = ((char) tiles[tn][i][j]) + 0x57;
      }
    }
    out[8*8] = '\0';
  }
}

// Returns number of used arrangements
int Tileset::getArrsCount() {
  int j, x, y;
  for (j = 1024; j > 0; j--)
    for (x = 0; x < 4; x++)
      for (y = 0; y < 4; y++)
        if (arrs[j - 1][x][y] != 0)
          return j;
   return 0;
}

void Tileset::getArrAsStr(int an, char* out) {
  int j,k;
  for (j=0; j<4; j++) {
    for (k=0; k<4; k++) {
      itoa(arrs[an][k][j] & 0xffff, out, 16);
      addZeros(out, 4);
      out += 4;
      itoa(collision[an][k][j] & 0xff, out, 16);
      addZeros(out, 2);
      out += 2;
    }
  }
  out = '\0';
}

void Tileset::setArrAsStr(int an, char* str) {
  int j,k;
  for (j=0; j<4; j++) {
    for (k=0; k<4; k++) {
      arrs[an][k][j] = atoi_b(str, 4, 16);
      str += 4;
      collision[an][k][j] = atoi_b(str, 2, 16);
      str += 2;
    }
  }
}

void Tileset::addPalette(int mtset, int mpal, uint start, Rom &rom) {
  Palette* p = new Palette(mtset, mpal);
  p->readFromRom(start, rom);
  pals.push_back(p);
}

void Tileset::addPalette(char* str) {
  Palette* p = new Palette(str);
  pals.push_back(p);
}

TileModule::~TileModule() {
  int i;
  for (i=0; i<NUM_TSETS; i++)
    tsets[i].clearPalettes();
} 

const char* TileModule::getName() {
  return "Tilesets";
}

const char* TileModule::getCredits() {
  return "Written by MrTenda\nBased on JHack Tile Editor";
}

void TileModule::readFromRom(Rom &rom) {
  int i,j,k,t,a,arrsLen;
  uint collisionAddr, tmp;
  
  uchar buffer[32768];
  for (i=0; i<32768; i++)
    buffer[i]=0;

  for (i = 0; i < NUM_TSETS; i++) {
    // Read minitiles
    if (rom.decomp(toRegPointer(rom.readMulti(TILES_PTRS + i*4, 4)), buffer, 28673) != 28673) {
      cerr << "Error: Bad compressed data on tileset " << i << endl;
    } else {
      for (j=0; j<NUM_TILES; j++)
        read4BppImage(&buffer[32*j], (uchar*) (tsets[i].tiles[j]), 0);
    }

    // Read arrangements
    if ((arrsLen = rom.decomp(toRegPointer(rom.readMulti(0x2F12AB + (i * 4), 4)), buffer, 32768)) < 0) {
      cerr << "Error: Bad compressed data on tileset " << i << " arrangements." << endl;
    } else {
      a=0;
      arrsLen /= 32;
      for (t=0;t<arrsLen/*1024*/;t++)
        for (j=0;j<4;j++)
          for(k=0;k<4;k++) {
            tsets[i].arrs[t][k][j] = (buffer[a] & 0xff) + ((buffer[a+1] & 0xff) << 8);
            a+=2;
          }
      for (; t<1024; t++)
        for (j=0;j<4;j++)
          for (k=0;k<4;k++)
            tsets[i].arrs[t][k][j] = 0;
    }

    // Read collision
    collisionAddr = toRegPointer(rom.readMulti(0x2F137B + (i * 4), 4));
    for (j=0; j<arrsLen; j++) {
      tmp = rom.readMulti(collisionAddr + (j*2), 2);
      for (k=0; k<16; k++)
        tsets[i].collision[j][k%4][k/4] = rom.read(tmp + 0x180200 + k);
    }
    for (; j<1024; j++)
      for (k=0; k<16; k++)
        tsets[i].collision[j][k%4][k/4] = 0;
  }

  // Read palettes
  for (i=0; i < NUM_MAP_TSETS; i++) {
    // read the "drawing" tileset #
    t = rom.read(0x2f121b + 2*i);
    // estimate the # of palettes for this map tileset
    if (i==NUM_MAP_TSETS-1)
      k = 0xDAFAA7 - rom.readMulti(0x2F12FB + (i * 4), 4);
    else
      k = rom.readMulti(0x2f12fb + ((i+1)*4),4) - rom.readMulti(0x2f12fb + (i*4),4);
    k /= 0xc0;

    for (j=0; j<k; j++) {
      tsets[t].addPalette(i, j, toRegPointer(rom.readMulti(0x2F12FB + (i * 4), 4)) + 0xC0 * j, rom);
    }
  }

  // Set alternate mtset/mpal stuff, can only do this after all palette data has been read
  list<Palette*> *pals;
  list<Palette*> *pals2;
  list<Palette*>::iterator it, it2;
  uint altAddr;
  for (i=0; i < NUM_TSETS; i++) {
    pals = tsets[i].palettesList();
    for (it = pals->begin(); it != pals->end(); it++) {
      if ((altAddr = (*it)->getAltAddr()) == 0)
        (*it)->setAltPal(-1,-1);
      else {
        for (j=0; j<NUM_TSETS; j++) {
          pals2 = tsets[j].palettesList();
          for (it2 = pals2->begin(); it2 != pals2->end(); it2++) {
            if (altAddr == (*it2)->getAddr()) {
              (*it)->setAltPal((*it2)->getMapTileset(), (*it2)->getMapPalette());
              break;
            }
          }
          if (it2 != pals2->end())
            break;
        }
      }
    }
  }
}

struct CollEntry {
  uchar* dataPtr;
  uint romPos;
};

void TileModule::writeToRom(Rom &rom) {
  uchar dataBuffer[32768];
  uchar compBuffer[65536];

  int i, j, k, m, a, compLen, arrsCount;

  uint collDataPos = 0, palPos = 0x1a7ea7;

  CollEntry* coll;
  list<CollEntry*> colls;
  list<CollEntry*>::iterator coll_it;

  for (i=0; i<32768; i++)
    dataBuffer[i] = 0;

  for (i=0; i < NUM_TSETS; i++) {
    // cout << "Writing tileset #" << i << endl;
    // Write minitiles
    for (j=0; j < NUM_TILES; j++)
      write4BppImage((uchar*) (tsets[i].tiles[j]), &dataBuffer[32*j], 0);
    compLen = rom.comp(dataBuffer, compBuffer, 28673);
    writeToFree(rom, compBuffer, compLen, TILES_PTRS + i*4, 4);

    // Write arrangements
    arrsCount = tsets[i].getArrsCount();
    a=0;
    for (j=0; j < arrsCount; j++) {
      for (k=0; k<4; k++)
        for (m=0; m<4; m++) {
          dataBuffer[a++] = tsets[i].arrs[j][m][k] & 0xff;
          dataBuffer[a++] = tsets[i].arrs[j][m][k] >> 8;
        }
    }
    compLen = rom.comp(dataBuffer, compBuffer, a);
    writeToFree(rom, compBuffer, compLen, 0x2F12AB + i*4, 4);

    // Write collisions
    for (j=0; j < arrsCount; j++) {
      // Search the list of coll data to see if there's an identical collision entry
      for (coll_it = colls.begin(); coll_it != colls.end(); coll_it++) {
         if (memcmp((*coll_it)->dataPtr, tsets[i].collision[j], 16*sizeof(uchar)) == 0) {
           // Found it
           dataBuffer[j*2] = (*coll_it)->romPos & 0xff;
           dataBuffer[j*2+1] = ((*coll_it)->romPos >> 8) & 0xff;
           break;
         }
      }
      if (coll_it == colls.end()) {
        // Need to add a new collission entry
        if (collDataPos + 16 >= 0x18f25e)
          cerr << "Error: Not enough room for collision data" << endl;
        else {
          dataBuffer[j*2] = collDataPos & 0xff;
          dataBuffer[j*2 + 1] = (collDataPos >> 8) & 0xff;
          for (k=0; k<16; k++) {
            rom.write(tsets[i].collision[j][k%4][k/4], 0x180200 + collDataPos + k);
          }
          coll = new CollEntry;
          coll->dataPtr = (uchar*) tsets[i].collision[j];
          coll->romPos = collDataPos;
          colls.push_back(coll);
          collDataPos += 16;
        }
      }
    }
    writeToFree(rom, dataBuffer, arrsCount * 2, 0x2F137B + (i*4), 4);
  }

  // Free collision list
  for (coll_it = colls.begin(); coll_it != colls.end(); coll_it++)
    delete (*coll_it);
  colls.clear();

  // Write palettes
  list<Palette*> *pals, *pals2;
  list<Palette*>::iterator it, it2;
  for (i=0; i < NUM_MAP_TSETS; i++) {
    // Set j to the drawing tileset #
    for (j=0; j < NUM_TSETS; j++) {
      pals = tsets[j].palettesList();
      for (it = pals->begin(); it != pals->end(); it++) {
        if ((*it)->getMapTileset() == i)
          break;
      }
      if (it != pals->end())
        break;
    }
    // Write the drawing tileset # to ROM
    rom.write(j, 0x2f121b + 2*i);
    // Write the pointer to the palette data for this map tileset
    rom.writeMulti(toSnesPointer(palPos), 0x2f12fb + i*4, 4);
    // Write the palettes to ROM
    j=0;
    for (; it != pals->end(); it++) {
      if ((*it)->getMapTileset() == i) {
        (*it)->writeToRom(palPos + 0xc0 * (*it)->getMapPalette(), rom);
        j++;
      }
    }
    palPos += j*0xc0;
  }

  // Write alt palette data, can only write this after we've written all the palettes
  for (i=0; i < NUM_TSETS; i++) {
    pals = tsets[i].palettesList();
    for (it = pals->begin(); it != pals->end(); it++) {
      if (((*it)->getFlag() != 0) && ((*it)->getAltTset() >= 0) && ((*it)->getAltPal() >= 0)) {
        for (j=0; j < NUM_TSETS; j++) {
          pals2 = tsets[j].palettesList();
          for (it2 = pals2->begin(); it2 != pals2->end(); it2++) {
            if (((*it)->getAltTset() == (*it2)->getMapTileset()) && ((*it)->getAltPal() == (*it2)->getMapPalette())) {
              rom.writeMulti((*it2)->getAddr(), 0x1A0200 + (*it)->getAddr() + 0x20, 2); 
              break;
            }
          }
          if (it2 != pals2->end())
            break;
        }
      } else {
        rom.writeMulti(0, 0x1A0200 + (*it)->getAddr() + 0x20, 2);
      }
    }
  }
}

// Exports to the JHack "fts" format
void TileModule::writeToDat(char* dirName, map<string, string> &fns) {
  int i,j;
  char buf[300], buf2[20], buf3[20];
  stringstream ss;
  ofstream tFile;

  list<Palette*>* pals;
  list<Palette*>::iterator it;

  TiXmlDocument doc;
  TiXmlDeclaration * decl = new TiXmlDeclaration( "1.0", "", "" );
  doc.LinkEndChild( decl );

  TiXmlElement* root = new TiXmlElement("Palettes");
  doc.LinkEndChild(root);
  TiXmlElement* chr;
  TiXmlElement* elem;

  for (i=0; i<NUM_TSETS; i++) {
    strcpy(buf2, "tileset_");
    if (i < 10)
      strcat(buf2, "0");
    ss.str(std::string());
    ss << i;
    strcat(buf2, ss.str().c_str());
    strcpy(buf3, buf2);
    strcat(buf3, ".fts");
    getDatFname(fns, buf2, buf3, dirName, buf);

    tFile.open(buf);
    for (j=0; j<512; j++) {
      tsets[i].getTileAsStr(j, buf);
      tFile << buf << "\n";
      tsets[i].getTileAsStr(j^512, buf);
      tFile << buf << "\n\n";
    }

    tFile << "\n";

    pals = tsets[i].palettesList();
    for (it = pals->begin(); it != pals->end(); it++) {
      (*it)->getAsString(buf);
      tFile << buf << "\n";

      // Special palette attributes in an external file
      chr = new TiXmlElement("Palette");
      chr->SetAttribute("map_tset", (*it)->getMapTileset());
      chr->SetAttribute("map_pal", (*it)->getMapPalette());
      root->LinkEndChild(chr);

      elem = new TiXmlElement("Flag");
      if ((*it)->getFlag() == 0) {
        elem->LinkEndChild( new TiXmlText( "0" ));
      } else {
        elem->SetAttribute("map_tset", (*it)->getAltTset());
        elem->SetAttribute("map_pal", (*it)->getAltPal());
        ss.str(std::string());
        ss << hex << (int) (*it)->getFlag() << dec;
        elem->LinkEndChild( new TiXmlText( ss.str().c_str() ));
      }
      chr->LinkEndChild(elem);

      elem = new TiXmlElement("SpritePalette");
      ss.str(std::string());
      ss << (int) (*it)->getSpritePalette();
      elem->LinkEndChild( new TiXmlText( ss.str().c_str() ));
      chr->LinkEndChild(elem);

      elem = new TiXmlElement("FlashEffect");
      ss.str(std::string());
      ss << (int) (*it)->getFlashEffect();
      elem->LinkEndChild( new TiXmlText( ss.str().c_str() ));
      chr->LinkEndChild(elem);
    }

    tFile << "\n";

    for (j=0; j<1024; j++) {
      tsets[i].getArrAsStr(j, buf);
      tFile << "\n" << buf;
    }

    tFile.close();
  }

  getDatFname(fns, "pal_attrs", "pal_attrs.xml", dirName, buf);
  doc.SaveFile(buf);
}

// reads from JHack "fts" format
void TileModule::readFromDat(char* dirName, map<string, string> &fns, map<string, int> &labels) {
  char buf[300], buf2[20], buf3[20];
  int i, tn, newlnCount, mode;
  stringstream ss;
  ifstream tFile;

  for (i=0; i<NUM_TSETS; i++) { 
    strcpy(buf2, "tileset_");
    if (i < 10)
      strcat(buf2, "0");
    ss.str(std::string());
    ss << i;
    strcat(buf2, ss.str().c_str());
    strcpy(buf3, buf2);
    strcat(buf3, ".fts");
    getDatFname(fns, buf2, buf3, dirName, buf);
    
    tFile.open(buf);
    tn = 0, newlnCount = 0;
    mode = 0;
    if (tFile.is_open()) {
      while (tFile.getline(buf,300)) {
        if (strlen(buf) == 0) {
          newlnCount++;
        } else {
          if (newlnCount >= 2) {
            mode++;
            tn = 0;
          }
          newlnCount = 0;

          if (mode == 0) { // Tileset mode
            tsets[i].setTileAsStr(tn, buf);
            tFile.getline(buf,300);
            if ((tn^512) < NUM_TILES)
              tsets[i].setTileAsStr(tn^512, buf);
            tn++;
          } else if (mode == 1) { // Palette mode
            tsets[i].addPalette(buf);
          } else if (mode == 2) { // Arrangement mode
            tsets[i].setArrAsStr(tn, buf);
            tn++;
          }
        }
      }
    }
    tFile.close();
  }

  // Read color0 data
  getDatFname(fns, "pal_attrs", "pal_attrs.xml", dirName, buf);
  TiXmlDocument doc(buf);
  if (!doc.LoadFile()) {
    cerr << "Error: (" << getName() << ") Failed to load " << buf << endl;
    return;
  }
  TiXmlHandle hDoc(&doc);
  TiXmlElement* elem;
  TiXmlHandle hRoot(0), hPal(0);

  // Get root element
  elem = hDoc.FirstChildElement().Element();
  hRoot = TiXmlHandle(elem);

  //Traverse Pals
  TiXmlElement* palNode = hRoot.FirstChild("Palette").Element();
  TiXmlElement* attrNode;
  const char* elemKey;
  int mt, mp;
  list<Palette*>* pals;
  list<Palette*>::iterator it;
  int tmp;
  for (; palNode; palNode = palNode->NextSiblingElement()) {
    mt = strToInt(palNode->Attribute("map_tset"));
    mp = strToInt(palNode->Attribute("map_pal"));

    // Find the matching entry
    for (i=0; i < NUM_TSETS; i++) {
      pals = tsets[i].palettesList();
      for (it = pals->begin(); it != pals->end(); it++) {
        if (((*it)->getMapTileset() == mt) && ((*it)->getMapPalette() == mp))
          break;
      }
      if (it != pals->end())
        break;
    }

    hPal = TiXmlHandle(palNode);
    attrNode = hPal.FirstChild().Element();
    for (; attrNode; attrNode = attrNode->NextSiblingElement()) {
      elemKey = attrNode->Value();
      if (strcmp(elemKey, "Flag") == 0) {
        tmp = strToInt(attrNode->GetText(), 16);
        if (tmp == 0)
          (*it)->setFlag(0);
        else {
          (*it)->setFlag(tmp);
          (*it)->setAltPal(strToInt(attrNode->Attribute("map_tset")), strToInt(attrNode->Attribute("map_pal")));
        }
      } else if (strcmp(elemKey, "SpritePalette") == 0)
        (*it)->setSpritePalette(strToInt(attrNode->GetText()));
      else if (strcmp(elemKey, "FlashEffect") == 0)
        (*it)->setFlashEffect(strToInt(attrNode->GetText()));
    }
  }
}
